import os
import copy
import torch
from tqdm import tqdm
from torch import optim, save
from torch.utils.data import DataLoader
from torchmetrics import ScaleInvariantSignalNoiseRatio

from utils import AudioDataset, plot_line
from model import TasNet

# 忽略一些警告输出
import warnings
warnings.filterwarnings("ignore")


# 尝试用GPU
def try_gpu(i=0):
    """如果存在 GPU(i) ，就返回，否则返回 cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')


def cal_loss(net, loss, data, device):
    """ 计算测试集数据的总损失
    """
    gross_loss = []
    net.eval()
    print("start test cal loss ...")
    for x, y in tqdm(data):
        x, y = x.to(device), y.to(device)
        y_hat = net(x)
        gross_loss.append(loss(y_hat, y).item())
    
    return gross_loss


if __name__ == "__main__":
    # 参数设置
    lr = 1e-3
    epochs = 3
    batch_size = 2

    # 准备数据加载器
    train = DataLoader(AudioDataset("./data/TIMIT-mix-8/TRAIN"), batch_size=batch_size, shuffle=True)
    test = DataLoader(AudioDataset("./data/TIMIT-mix-8/TEST"), batch_size=batch_size)

    # 训练准备
    device = try_gpu()
    print("device: ", device)
    net = TasNet().to(device)
    optimizer = optim.Adam(net.parameters(), lr=lr)
    lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=5)
    loss = ScaleInvariantSignalNoiseRatio().to(device)
    best = {'loss': float('inf'), 'net': None, 'epoch': None}

    # 开始训练
    loss_table = []
    for epoch in range(epochs):
        net.train()
        print(f"epoch {epoch + 1} start train:")
        gross_loss = 0.0
        for x, y in tqdm(train):
            x, y = x.to(device), y.to(device)
            y_hat = net(x)
            los = loss(y_hat, y)
            l = los.item()
            loss_table.append(l)
            gross_loss += l

            optimizer.zero_grad()
            los.backward()
            optimizer.step()
            lr_scheduler.step(los)

        test_loss = cal_loss(net, loss, test, device)
        plot_line(test_loss, f"epoch{epoch+1}-test-loss")
        test_loss = sum(test_loss)
        print("train loss: {:.5f}  test loss: {:.5f}".format(gross_loss, test_loss))

        if test_loss < best['loss']:
            best['loss'] = test_loss
            best['net'] = copy.deepcopy(net)
            best['epoch'] = epoch + 1
    
    # 绘制训练过程 loss 值的变化曲线
    plot_line(loss_table, f"train-loss")

    # 保存模型
    save(best['net'].state_dict(), os.path.join("./save_net", f"epoch{best['epoch']}-Conv-TasNet.pt"))
    print(f"test gross loss of save net \'epoch{best['epoch']}-Conv-TasNet.pt\' is {best['loss']}")
